#-*- coding: utf-8 -*-
import os
import numpy as np
#高斯类，用于写高斯文件的写入，与高斯文件的输出。
class gauss():
#设置基本的高斯信息：
#name->文件名；spin->分子系统自旋；fragCharge->分子系统电荷量。tilte->分子系统名字。
#proc->使用cpu数；mem->使用内存数。
#我预先将每一个字段写成字符串（带换行符），最后在write时依次写入。因为key字段与之后
#分子系统后面需加入的字段有关，所以没办法只能单独写一个函数来设置。
#需要注意的是我会在每个大的字段（需要空行）后面再加一个换行符，以便空行。
  def setBasic(self,name,spin,fragCharge,fragments,title='X',proc='8',mem='600MW'):
    self.chk = '%chk=' + name + '.chk\n'
    self.log = name+'.log'
    self.name = name + '.gjf'
    self.process = '%nproc='+proc+'\n'
    self.memory = '%mem='+mem+'\n'
    self.title = 'EE-GMFCC AND MOLECULAR '+title+'\n\n'
    self.spinCharge = str(fragCharge)+' '+str(spin)+'\n'
#设定分子系统，fragments由多个碎片组成，比如残基+帽子，另外比如算两体时，分子系统都不只一个碎片，
#每个碎片都是一个元组（i,j)i是分子系统的坐标，j是分子系统的元素名。
    fragment = []
    for frag in fragments:
      atomName = frag[1]
      atomCrd = frag[0]
      for i in xrange(len(atomName)):
        fragment.append(' {0:<8}{1:>12.7f}{2:12.7f}{3:12.7f}\n'.format(atomName[i],atomCrd[i][0],\
                   atomCrd[i][1],atomCrd[i][2]))
    fragment[-1] = fragment[-1]+'\n'
#这里会将分子系统写成一个字符串，这样后面就可以直接写入文件。
    self.fragment = ''.join(fragment)
#设置关键字字段，这里设置下计算方法,基组，以及其他需要的关键字。（比如charge)
#然后在other字段里填上相应关键字设置后在分子系统后面需要添加的信息。(比如背景电荷）。
  def setKey(self,method='hf',basis='3-21g'):
    self.key = '#p force nosymm ' + method + '/' + basis + '\n\n'
    self.others = ''
#写高斯文件。
  def write(self):
    gjf = open(self.name,'w')
    gjf.write(self.chk)        #写入chk文件名。
    gjf.write(self.process)    #写入cpu数。
    gjf.write(self.memory)     #写入内存数。
    gjf.write(self.key)        #写入关键字。
    gjf.write(self.title)      #写入标题。
    gjf.write(self.spinCharge) #写入自旋与电荷。
    gjf.write(self.fragment)   #写入分子系统。
    gjf.write(self.others)     #写入其他相关信息。
    gjf.close()
#跑高斯
  def run(self):
    command = 'g09 '+self.name
    os.system(command)
#读取高斯输出文件。    
#设置name主要是为了之后检查时，不需要在计算一遍，可以直接读计算出的log文件。
  def read(self,name=None):
    if name == None:
      log = open(self.log,'r')
    else:
      log = open(name,'r')
    records = []
    while 1:
      line = log.readline()
      if 'SCF Done:' in line: #读能量。
        self.calcEnergy = float(line[23:39])
      elif 'Forces (Hartrees/Bohr)' in line: #读力。
        log.readline();log.readline()
        while 1:
          line = log.readline()
          if '---' in line:
            break
          records.append(line[24:])
      if line == '':
        break
    records = ''.join(records)
    forces = np.fromstring(records, dtype=float, sep=' ')
    self.calcForce = forces.reshape(forces.size/3,3)
     
#加背景电荷的高斯
class backChargeGauss(gauss):
#重写后的方法会读取电荷系统，并在key中加入charge。
  def setKey(self,charge,method='hf',basis='3-21g'):
    self.key = '#p force nosymm charge ' + method + '/' + basis + '\n\n'
    back = []
    backCrd = charge[0]
    backCharge = charge[1]
    for i in xrange(len(backCharge)):
      back.append(' {0:14.8f}{1:14.8f}{2:14.8f}{3:14.8f}\n'.format(backCrd[i][0],backCrd[i][1],\
                 backCrd[i][2],backCharge[i]))
    back[-1] += '\n'
#这里将other设置为背景电荷系统。
    self.others = ''.join(back)
#多读取背景电荷的能量。
  def read(self,name):
    log = open(name,'r')
    records = []
    while 1:
      line = log.readline()
      if 'Self energy of the charges' in line: #读取背景电荷能量。
        self.CalcChargeEnergy = float(line[30:50])
      elif 'SCF Done:' in line:
        self.CalcEnergy = float(line[23:39])
      elif 'Forces (Hartrees/Bohr)' in line:
        log.readline();log.readline()
        while 1:
          line = log.readline()
          if '---' in line:
            break
          records.append(line[24:])
      if line == '':
        break
    records = ''.join(records)
    forces = np.fromstring(records, dtype=float, sep=' ')
    self.CalcForce = forces.reshape(forces.size/3,3)

if __name__ == '__main__':
    a = gauss()
    spin = 1 ; charge = 1
    name = 'aaa'
    atom = [[1,1,1],[2,2,2],[3,3,3]]
    atomName = ['H','O','H']
    cube = (atom,atomName)
    a.setBasic(name,spin,charge,cube)
    a.setKey()
    q.write()

